package main

import (
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"regexp"
	"strings"
)

// GetImageURLs 从给定的URL获取页面内容，并提取所有img标签的src属性值，去重后返回
func GetImageURLs(url string) ([]string, error) {
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	// 由于 ioutil.ReadAll 已弃用，使用 io.ReadAll 替代
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	// 使用正则表达式匹配img标签的src属性
	re := regexp.MustCompile(`<img[^>]+src="([^">]+)"`)
	matches := re.FindAllStringSubmatch(string(body), -1)

	// 去重
	urlMap := make(map[string]bool)
	var uniqueURLs []string
	for _, match := range matches {
		if len(match) > 1 {
			imgURL := match[1]
			if !urlMap[imgURL] {
				urlMap[imgURL] = true
				uniqueURLs = append(uniqueURLs, imgURL)
			}
		}
	}

	return uniqueURLs, nil
}

// SaveURLsToFile 将URL列表保存到文件中

// 定义一个函数，添加一个参数用于控制是否追加内容
// SaveURLsToFile 支持将文本或 URL 切片保存到文件，可选择追加模式
func SaveURLsToFile(input interface{}, filePath string, appendMode bool) error {
	var data []byte
	switch v := input.(type) {
	case string:
		data = []byte(v + "\n")
	case []string:
		var content strings.Builder
		for _, url := range v {
			content.WriteString(url + "\n")
		}
		data = []byte(content.String())
	default:
		return fmt.Errorf("不支持的输入类型，仅支持 string 或 []string")
	}

	// 打开文件时，根据 appendMode 参数选择追加模式或覆盖模式
	var file *os.File
	var err error
	if appendMode {
		file, err = os.OpenFile(filePath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	} else {
		file, err = os.Create(filePath)
		if err != nil {
			return err
		}
	}
	defer file.Close()
	_, err = file.Write(data)
	return err

}

// GetAllURLs 从给定的URL获取页面内容，并提取所有a标签的href属性值，去重后返回
func GetAllURLs(urlStr string, maxDepth int) ([]string, error) {
	visited := make(map[string]bool)
	return getAllURLsRecursive(urlStr, maxDepth, visited)
}

func getAllURLsRecursive(urlStr string, depth int, visited map[string]bool) ([]string, error) {
	if depth <= 0 {
		return nil, nil
	}
	if visited[urlStr] {
		return nil, nil
	}
	visited[urlStr] = true

	resp, err := http.Get(urlStr)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	// 由于 ioutil.ReadAll 已弃用，使用 io.ReadAll 替代
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	// 使用正则表达式匹配a标签的href属性
	re := regexp.MustCompile(`<a[^>]+href="([^">]+)"`)
	matches := re.FindAllStringSubmatch(string(body), -1)
	// 解析当前网站域名
	// 原代码中 url 是字符串类型，没有 Parse 方法，需要导入 net/url 包并使用该包的 Parse 函数
	u, err := url.Parse(urlStr)
	if err != nil {
		return nil, err
	}
	baseURL := u.Scheme + "://" + u.Host

	// 处理匹配结果，添加域名
	for i, match := range matches {
		if len(match) > 1 {
			if !strings.HasPrefix(match[1], "http") {
				if !strings.HasPrefix(match[1], "/") {
					match[1] = "/" + match[1]
				}
				matches[i][1] = baseURL + match[1]
			}
		}
	}
	// 去重
	urlMap := make(map[string]bool)
	var uniqueURLs []string
	for _, match := range matches {
		// 过滤掉非当前网站域名的URL
		if !strings.HasPrefix(match[1], baseURL) {
			continue
		}
		if len(match) > 1 {
			linkURL := match[1]
			if !urlMap[linkURL] {
				urlMap[linkURL] = true
				uniqueURLs = append(uniqueURLs, linkURL)
			}
		}
	}
	// 递归抓取子链接
	var childURLs []string
	for _, link := range uniqueURLs {
		urls, err := getAllURLsRecursive(link, depth-1, visited)
		if err != nil {
			return nil, err
		}
		childURLs = append(childURLs, urls...)
	}
	newURLs := append(uniqueURLs, childURLs...)
	// err2 := os.WriteFile("./site_url_temp.txt", []byte(strings.Join(newURLs, "\n")), 0644)
	// if err2 != nil {
	// 	fmt.Printf("保存URL到文件时出错: %v\n", err)
	// }
	return newURLs, nil
}

func collectUrls(pageURL string) {
	// 替换为你要获取图片URL的网页地址
	if pageURL == "" {
		pageURL = "https://www.soogif.com/"
	}
	outputFilePath := "site_urls.txt"

	// 设置抓取深度为2（可根据需要调整）
	urlStr, urlStrErr := GetAllURLs(pageURL, 2)
	if urlStrErr != nil {
		fmt.Printf("获取图片URL时出错: %v\n", urlStrErr)
		return
	}
	err := os.WriteFile(outputFilePath, []byte(strings.Join(urlStr, "\n")), 0644)
	if err != nil {
		fmt.Printf("保存URL到文件时出错: %v\n", err)
		return
	}
}
func collectImages(pageURL string) {
	// 读取 site_urls.txt 文件内容
	fileContent, err := os.ReadFile("site_urls.txt")
	if err != nil {
		fmt.Printf("读取 site_urls.txt 文件时出错: %v\n", err)
		return
	}
	// 将文件内容按行分割成 URL 列表
	urlList := strings.Split(string(fileContent), "\n")

	// 循环调用 GetImageURLs 函数
	for _, url := range urlList {
		if url == "" {
			continue
		}
		imageURLs, err := GetImageURLs(url)
		if err != nil {
			fmt.Printf("从 %s 获取图片 URL 时出错: %v\n", url, err)
			continue
		}
		// 这里可以添加保存图片 URL 的逻辑，例如保存到文件
		outputFilePath := "image_urls_from_site.txt"
		SaveURLsToFile(imageURLs, outputFilePath, true)
	}
}

func urlClean(filename string) ([]string, error) {
	// 读取 site_urls.txt 文件内容
	fileContent, err := os.ReadFile(filename)
	if err != nil {
		fmt.Printf("读取 site_urls.txt 文件时出错: %v\n", err)
		return []string{}, err
	}
	// 将文件内容按行分割成 URL 列表
	urlList := strings.Split(string(fileContent), "\n")

	// 用于去重的 map
	urlMap := make(map[string]bool)
	var validURLs []string

	// 遍历 URL 列表，去掉非 http/https 协议的 URL 并去重
	for _, url := range urlList {
		if url == "" {
			continue
		}
		if strings.HasPrefix(url, "http://") || strings.HasPrefix(url, "https://") {
			if !urlMap[url] {
				urlMap[url] = true
				validURLs = append(validURLs, url)
			}
		}
	}
	return validURLs, err
}

func main() {
	filename := "image_urls_from_site.txt"
	newfilename := "image_urls_from_site_new.txt"
	urls, err := urlClean(filename)
	if err != nil {
		fmt.Printf("读取 site_urls.txt 文件时出错: %v\n", err)
		return
	}
	SaveURLsToFile(urls, newfilename, false)
	// // 替换为你要保存URL的文件路径
	// outputFilePath := "image_urls.txt"

	// urls, err := GetImageURLs(pageURL)
	// if err != nil {
	// 	fmt.Printf("获取图片URL时出错: %v\n", err)
	// 	return
	// }

	// err = SaveURLsToFile(urls, outputFilePath)
	// if err != nil {
	// 	fmt.Printf("保存URL到文件时出错: %v\n", err)
	// 	return
	// }

	// fmt.Println("图片URL已成功保存到文件")
}
